{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZwZNOAMZcxl3"
      },
      "source": [
        "##### Copyright 2022 The TensorFlow Neural Structured Learning Authors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nxbcnXODdE06"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-BszoQj0dSZO"
      },
      "source": [
        "# Saving and loading NSL models\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfqlePz0g6o5"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/saving_and_loading_nsl_models.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/saving_and_loading_nsl_models.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oHEGl8h_m6tS"
      },
      "source": [
        "## Overview\n",
        "\n",
        "In this example, we will demonstrate how to save and load `AdversarialRegularization` models in NSL. For more details about the concept\n",
        "of adversarial regularization and how to use the `AdversarialRegularization`\n",
        "API, please see this \n",
        "[tutorial](https://www.tensorflow.org/neural_structured_learning/tutorials/adversarial_keras_cnn_mnist)\n",
        "or this [video explanation](https://youtu.be/Js2WJkhdU7k).\n",
        "\n",
        "The high-level structure of this example is as follows:\n",
        "\n",
        "1.  Create a neural network as a base model.\n",
        "    * Here we create the base model with the `tf.keras` functional API.\n",
        "    * This procedure is also compatible with models created by the\n",
        "      `tf.keras.Sequential` API and models created by subclasses of\n",
        "      `tf.keras.Model`, provided that the subclass models support saving and loading.\n",
        "    * For more information on Keras models in TensorFlow, see this\n",
        "      [documentation](https://www.tensorflow.org/api_docs/python/tf/keras/Model).\n",
        "2.  Wrap the base model with the **`AdversarialRegularization`** wrapper class.\n",
        "    * The wrapper class is provided by the NSL framework, and it creates a new\n",
        "      `tf.keras.Model` instance.\n",
        "    * This new model is a subclass model and will include the adversarial loss\n",
        "      as a regularization term in its training objective.\n",
        "3.  Train and evaluate the new model.\n",
        "4.  Save the model, and load it back.\n",
        "5.  Verify that the loaded model is the same as the one before saving."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qODwGDl-n42t"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4RhmgQ7-mlrl"
      },
      "source": [
        "Install the Neural Structured Learning package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ByJ7133BQULR"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet neural-structured-learning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PZvsEQrhSqKx"
      },
      "source": [
        "Import libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EuqEuAYzTMo0"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import neural_structured_learning as nsl\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3LwBtQGaTvbe"
      },
      "source": [
        "## Hyperparameters\n",
        "\n",
        "We collect and explain the hyperparameters (in an `HParams` object) for model\n",
        "training and evaluation.\n",
        "\n",
        "Input/Output:\n",
        "\n",
        "*   **`input_shape`**: The shape of the input tensor. Each image is 28-by-28\n",
        "pixels with 1 channel.\n",
        "*   **`num_classes`**: There are a total of 10 classes, corresponding to 10\n",
        "digits [0-9].\n",
        "\n",
        "Model architecture:\n",
        "\n",
        "*   **`conv_filters`**: A list of numbers, each specifying the number of\n",
        "filters in a convolutional layer.\n",
        "*   **`kernel_size`**: The size of 2D convolution window, shared by all\n",
        "convolutional layers.\n",
        "*   **`pool_size`**: Factors to downscale the image in each max-pooling layer.\n",
        "*   **`num_fc_units`**: The number of units (i.e., width) of each\n",
        "fully-connected layer.\n",
        "\n",
        "Training and evaluation:\n",
        "\n",
        "*  **`batch_size`**: Batch size used for training and evaluation.\n",
        "*  **`epochs`**: The number of training epochs.\n",
        "\n",
        "Adversarial learning:\n",
        "\n",
        "*   **`adv_multiplier`**: The weight of adversarial loss in the training\n",
        "objective, relative to the labeled loss.\n",
        "*   **`adv_step_size`**: The magnitude of adversarial perturbation.\n",
        "*  **`adv_grad_norm`**: The norm to measure the magnitude of adversarial\n",
        "perturbation.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iOc8YdmIRSHo"
      },
      "outputs": [],
      "source": [
        "class HParams(object):\n",
        "  def __init__(self):\n",
        "    self.input_shape = [28, 28, 1]\n",
        "    self.num_classes = 10\n",
        "    self.conv_filters = [32, 64, 64]\n",
        "    self.kernel_size = (3, 3)\n",
        "    self.pool_size = (2, 2)\n",
        "    self.num_fc_units = [64]\n",
        "    self.batch_size = 32\n",
        "    self.epochs = 5\n",
        "    self.adv_multiplier = 0.2\n",
        "    self.adv_step_size = 0.2\n",
        "    self.adv_grad_norm = 'infinity'\n",
        "\n",
        "HPARAMS = HParams()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "72zL1AMcYYGG"
      },
      "source": [
        "## MNIST dataset\n",
        "\n",
        "The [MNIST dataset](http://yann.lecun.com/exdb/mnist/) contains grayscale\n",
        "images of handwritten digits (from '0' to '9'). Each image shows one digit at\n",
        "low resolution (28-by-28 pixels). The task involved is to classify images into\n",
        "10 categories, one per digit.\n",
        "\n",
        "Here we load the MNIST dataset from\n",
        "[TensorFlow Datasets](https://www.tensorflow.org/datasets). It handles\n",
        "downloading the data and constructing a `tf.data.Dataset`. The loaded dataset\n",
        "has two subsets:\n",
        "\n",
        "*   `train` with 60,000 examples, and\n",
        "*   `test` with 10,000 examples.\n",
        "\n",
        "Examples in both subsets are stored in feature dictionaries with the following\n",
        "two keys:\n",
        "\n",
        "*   `image`: Array of pixel values, ranging from 0 to 255.\n",
        "*   `label`: Groundtruth label, ranging from 0 to 9."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R1dK6E4axNHB"
      },
      "outputs": [],
      "source": [
        "datasets = tfds.load('mnist')\n",
        "\n",
        "train_dataset = datasets['train']\n",
        "test_dataset = datasets['test']\n",
        "\n",
        "IMAGE_INPUT_NAME = 'image'\n",
        "LABEL_INPUT_NAME = 'label'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IBkh4mbsxLR_"
      },
      "source": [
        "To make the model numerically stable, we normalize the pixel values to [0, 1]\n",
        "by mapping the dataset over the `normalize` function. After shuffling training\n",
        "set and batching, we convert the examples to feature tuples `(image, label)`\n",
        "for training the base model. We also provide a function to convert from tuples\n",
        "to dictionaries for later use."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VhMEJqKs0_7z"
      },
      "outputs": [],
      "source": [
        "def normalize(features):\n",
        "  features[IMAGE_INPUT_NAME] = tf.cast(\n",
        "      features[IMAGE_INPUT_NAME], dtype=tf.float32) / 255.0\n",
        "  return features\n",
        "\n",
        "def convert_to_tuples(features):\n",
        "  return features[IMAGE_INPUT_NAME], features[LABEL_INPUT_NAME]\n",
        "\n",
        "def convert_to_dictionaries(image, label):\n",
        "  return {IMAGE_INPUT_NAME: image, LABEL_INPUT_NAME: label}\n",
        "\n",
        "train_dataset = train_dataset.map(normalize).shuffle(10000).batch(HPARAMS.batch_size).map(convert_to_tuples)\n",
        "test_dataset = test_dataset.map(normalize).batch(HPARAMS.batch_size).map(convert_to_tuples)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JrrMpPNmpCKK"
      },
      "source": [
        "## Base model\n",
        "\n",
        "Our base model will be a neural network consisting of 3 convolutional layers\n",
        "follwed by 2 fully-connected layers (as defined in `HPARAMS`). Here we define\n",
        "it using the Keras functional API. Feel free to try other APIs or model\n",
        "architectures (e.g. subclassing). Note that the NSL framework does support all three types of Keras APIs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4UjrtuIsYWo3"
      },
      "outputs": [],
      "source": [
        "def build_base_model(hparams):\n",
        "  \"\"\"Builds a model according to the architecture defined in `hparams`.\"\"\"\n",
        "  inputs = tf.keras.Input(\n",
        "      shape=hparams.input_shape, dtype=tf.float32, name=IMAGE_INPUT_NAME)\n",
        "\n",
        "  x = inputs\n",
        "  for i, num_filters in enumerate(hparams.conv_filters):\n",
        "    x = tf.keras.layers.Conv2D(\n",
        "        num_filters, hparams.kernel_size, activation='relu')(\n",
        "            x)\n",
        "    if i < len(hparams.conv_filters) - 1:\n",
        "      # max pooling between convolutional layers\n",
        "      x = tf.keras.layers.MaxPooling2D(hparams.pool_size)(x)\n",
        "  x = tf.keras.layers.Flatten()(x)\n",
        "  for num_units in hparams.num_fc_units:\n",
        "    x = tf.keras.layers.Dense(num_units, activation='relu')(x)\n",
        "  pred = tf.keras.layers.Dense(hparams.num_classes)(x)\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=pred)\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "288nsmN5pLoo"
      },
      "outputs": [],
      "source": [
        "base_model = build_base_model(HPARAMS)\n",
        "base_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CemXA8N9q336"
      },
      "source": [
        "## Adversarial-regularized model\n",
        "\n",
        "Here we show how to incorporate adversarial training into a Keras model with a\n",
        "few lines of code, using the NSL framework. The base model is wrapped to create\n",
        "a new `tf.Keras.Model`, whose training objective includes adversarial\n",
        "regularization."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YUOpl-rkzRrY"
      },
      "source": [
        "First, we create a config object with all relevant hyperparameters using the\n",
        "helper function `nsl.configs.make_adv_reg_config`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-WWVwJB2qstE"
      },
      "outputs": [],
      "source": [
        "adv_config = nsl.configs.make_adv_reg_config(\n",
        "    multiplier=HPARAMS.adv_multiplier,\n",
        "    adv_step_size=HPARAMS.adv_step_size,\n",
        "    adv_grad_norm=HPARAMS.adv_grad_norm\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OmeIUyxE4s68"
      },
      "source": [
        "Now we can wrap a base model with `AdversarialRegularization`. Here we create a\n",
        "new base model (`base_adv_model`), so that the existing one (`base_model`) can\n",
        "be used in later comparison.\n",
        "\n",
        "The returned `adv_model` is a `tf.keras.Model` object, whose training objective\n",
        "includes a regularization term for the adversarial loss. To compute that loss,\n",
        "the model has to have access to the label information (feature `label`), in\n",
        "addition to regular input (feature `image`). For this reason, we convert the\n",
        "examples in the datasets from tuples back to dictionaries. And we tell the\n",
        "model which feature contains the label information via the `label_keys`\n",
        "parameter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TObqJLEX4sQq"
      },
      "outputs": [],
      "source": [
        "base_adv_model = build_base_model(HPARAMS)\n",
        "adv_model = nsl.keras.AdversarialRegularization(\n",
        "    base_adv_model,\n",
        "    label_keys=[LABEL_INPUT_NAME],\n",
        "    adv_config=adv_config\n",
        ")\n",
        "\n",
        "train_set_for_adv_model = train_dataset.map(convert_to_dictionaries)\n",
        "test_set_for_adv_model = test_dataset.map(convert_to_dictionaries)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sdnq7_7hNbDi"
      },
      "outputs": [],
      "source": [
        "adv_model = nsl.keras.AdversarialRegularization(\n",
        "    base_adv_model,\n",
        "    label_keys=[LABEL_INPUT_NAME],\n",
        "    adv_config=adv_config\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aKTQWzfj7JvL"
      },
      "source": [
        "Next we compile, train, and evaluate the\n",
        "adversarial-regularized model. There might be warnings like\n",
        "\"Output missing from loss dictionary,\" which is fine because\n",
        "the `adv_model` doesn't rely on the base implementation to\n",
        "calculate the total loss.\n",
        "\n",
        "Along training, we save the model weights in a checkpoint after each epoch.\n",
        "The checkpoint can be restored by using `adv_model.load_weights`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yk2tVJbZToIt"
      },
      "outputs": [],
      "source": [
        "adv_model.compile(\n",
        "    optimizer='adam',\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])\n",
        "\n",
        "adv_model.fit(\n",
        "    train_set_for_adv_model,\n",
        "    epochs=2,\n",
        "    callbacks=[tf.keras.callbacks.ModelCheckpoint('ckpt')])"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "results = adv_model.evaluate(test_set_for_adv_model)\n",
        "for metric, value in zip(adv_model.metrics_names, results):\n",
        "  print(metric, value)"
      ],
      "metadata": {
        "id": "OqHT07mTiFFk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Save and Load\n",
        "\n",
        "After training, we can save the model in a SavedModel format, which stores both\n",
        "architecture and weights in the model. For `AdversarialRegularization` models,\n",
        "the `save` method saves the base model because the `AdversarialRegularization`\n",
        "wrapper doesn't change the architecture nor add any variables in model weights."
      ],
      "metadata": {
        "id": "107Qc3SAzA8W"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lr2DVk9WKLD3"
      },
      "outputs": [],
      "source": [
        "adv_model.save('base_model') # only save base model"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can load the saved model using `tf.keras.models.load_model`. The model\n",
        "summary shows the same architecture as we saw in the [Base model](#scrollTo=JrrMpPNmpCKK) section."
      ],
      "metadata": {
        "id": "JmY9dR7-yIRS"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8gHOcOY2Kg78"
      },
      "outputs": [],
      "source": [
        "loaded_base_model = tf.keras.models.load_model('base_model')\n",
        "loaded_base_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can also verify that the model weights are the same."
      ],
      "metadata": {
        "id": "bg5YeFWly_Pw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "original_weights = adv_model.base_model.trainable_weights\n",
        "loaded_weights = loaded_base_model.trainable_weights\n",
        "\n",
        "all(np.array_equal(original_weight.numpy(), loaded_weight.numpy())\n",
        "    for original_weight, loaded_weight in zip(original_weights, loaded_weights))"
      ],
      "metadata": {
        "id": "mvc5lmGildhU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And so are the predictions and evaluation metrics."
      ],
      "metadata": {
        "id": "-iSIsKvCzjUi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "adv_base_preds = adv_model.base_model.predict(test_dataset)\n",
        "load_base_preds = loaded_base_model.predict(test_dataset)\n",
        "\n",
        "np.array_equal(adv_base_preds, load_base_preds)"
      ],
      "metadata": {
        "id": "eVbIfTHpmXl4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "adv_model.base_model.evaluate(test_dataset)"
      ],
      "metadata": {
        "id": "Sjw_6KJr1wQx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loaded_base_model.evaluate(test_dataset)"
      ],
      "metadata": {
        "id": "Noyn3l57nhGf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you would like to continue adversarial training on the loaded model, you can\n",
        "wrap it with `AdversarialRegularization` again."
      ],
      "metadata": {
        "id": "BsRa1uDp4Qcb"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3McCKw4hIe42"
      },
      "outputs": [],
      "source": [
        "# We need to set the adv_config and other related parameters again.\n",
        "adv_loaded_model = nsl.keras.AdversarialRegularization(\n",
        "    loaded_base_model,\n",
        "    label_keys=[LABEL_INPUT_NAME],\n",
        "    adv_config=adv_config\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w7njatc3TXhm"
      },
      "outputs": [],
      "source": [
        "adv_loaded_model.compile(\n",
        "    optimizer='adam',\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])\n",
        "adv_loaded_model.fit(\n",
        "    train_set_for_adv_model,\n",
        "    epochs=3,\n",
        "    callbacks=[tf.keras.callbacks.ModelCheckpoint('ckpt')])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BFj6YvUVTZXc"
      },
      "outputs": [],
      "source": [
        "results = adv_loaded_model.evaluate(test_set_for_adv_model)\n",
        "for metric, value in zip(adv_loaded_model.metrics_names, results):\n",
        "  print(metric, value)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can see that the continued training improved the model."
      ],
      "metadata": {
        "id": "3prVCBjt5vQk"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Saving and loading NSL models",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.8"
    },
    "toc": {
      "base_numbering": 1,
      "nav_menu": {},
      "number_sections": true,
      "sideBar": true,
      "skip_h1_title": false,
      "title_cell": "Table of Contents",
      "title_sidebar": "Contents",
      "toc_cell": false,
      "toc_position": {},
      "toc_section_display": true,
      "toc_window_display": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
